# encoding: UTF-8

from vnpy.event import EventEngine2 as EventEngine
from vnpy.trader.gateway import ctpGateway
from vnpy.trader.vtEngine import *
from vnpy.trader.uiQt import createQApp

target_symbol = 'cu1807'

predicted_buy_price = 54100
predicted_sell_price = 54200


class Main:
    def __init__(self):
        self.fully_initialized = False
        self.price_initialized = False
        self.subscribed = False
        self.timer_count = 0

        self.event_engine = EventEngine()
        self.event_engine.register(EVENT_TIMER, self.on_timer)
        self.event_engine.register(EVENT_TICK, self.on_tick)

        self.main_engine = MainEngine(eventEngine=self.event_engine)
        self.main_engine.addGateway(gatewayModule=ctpGateway)

        # 为了在非交易日也能获取行情，我修改了一下调用的先后顺序
        # 这句话在调到on_timer里面去了。
        # self.subscribe_market_data()
        pass

    def subscribe_market_data(self):
        """订阅我们感兴趣的合约的行情"""
        subscribe_req = VtSubscribeReq()
        subscribe_req.symbol = target_symbol  # 合约代码
        # subscribe_req.exchange = constant.EMPTY_STRING    # 交易所代码
        self.main_engine.subscribe(subscribe_req, ctpGateway.gatewayName)

    def current_long_position(self, symbol):
        """返回买的仓位"""
        for data in self.main_engine.getAllPositions():  # type: VtPositionData
            if data.symbol == symbol and data.direction == constant.DIRECTION_LONG:
                return data.position
        return 0

    def last_price(self, symbol):
        """获取某个合约的当前价格"""
        tick = self.main_engine.getTick(symbol)  # type: VtTickData
        return tick.lastPrice if tick else None

    def do_strategy(self):
        """很简单的策略：高于特定价格就卖，低于特定价格就买"""
        current_position = self.current_long_position(target_symbol)
        price = self.last_price(target_symbol)
        if price:
            if price < predicted_buy_price:
                if current_position == 0:
                    self.buy_one()
            else:  # price >= predicted_sell_price:
                if current_position > 0:
                    self.sell_all()

    def buy_one(self):
        """买开我们感兴趣的合约一手"""
        order = VtOrderReq()
        order.symbol = target_symbol  # 合约代码
        # order.exchange = constant.EMPTY_STRING          # 交易所, ctp中不需要
        # order.vtSymbol = symbol                         # VT合约代码, ctp中不需要
        order.price = 0  # 价格
        order.volume = 1  # 数量
        order.direction = constant.DIRECTION_LONG  # 买/卖(多头、空头)
        order.priceType = constant.PRICETYPE_MARKETPRICE  # 价格类型
        order.offset = constant.OFFSET_OPEN  # 开平
        # order.currency = currency                       # 货币类型, ctp中不需要
        # order.productClass = productClass               # 合约类型, ctp中不需要
        self.main_engine.sendOrder(order, ctpGateway.gatewayName)

    def sell(self, symbol, direction, volume):
        """平仓"""
        order = VtOrderReq()
        order.symbol = symbol
        order.price = 0
        order.volume = volume
        order.direction = direction
        order.priceType = constant.PRICETYPE_MARKETPRICE
        order.offset = constant.OFFSET_CLOSE  # 平仓
        self.main_engine.sendOrder(order, ctpGateway.gatewayName)

    def cancel_all_working_orders(self):
        """撤掉所有挂单"""
        for order in self.main_engine.getAllWorkingOrders():  # type: VtOrderData
            self.main_engine.cancelOrder(order, order.gatewayName)
        pass

    def sell_all_volume(self):
        """平掉所有持仓"""
        for position in self.main_engine.getAllPositions():  # type: VtPositionData
            self.sell(position.symbol, position.direction, position.position)

    def sell_all(self):
        """撤掉所有挂单、平掉所有持仓"""
        self.cancel_all_working_orders()
        self.sell_all_volume()

    def start(self):
        self.main_engine.connect(ctpGateway.gatewayName)

    def on_timer(self, event):
        # 如果还没订阅行情，则订阅行情。
        # 在这里订阅行情，即使在非交易日也能获取一个行情消息
        if not self.subscribed:
            self.subscribe_market_data()
            self.subscribed = True

        # 等待10s，让DataEngine（MainEngine的一部分）缓存一些数据，例如仓位、资金之类的
        if not self.fully_initialized:
            self.timer_count += 1
            if self.timer_count > 10:
                self.fully_initialized = True
            return

        # 等完10s，再开始我们的策略
        self.do_strategy()
        pass

    def on_tick(self, event):
        pass


def main():
    app = createQApp()  # 初始化Qt环境
    m = Main()  # 初始化我们的策略程序
    m.start()  # 连接到服务器
    exit(app.exec_())  # 让程序自己跑


if __name__ == '__main__':
    main()
